#include "swarm_computer.hpp"
#include <iostream>
#include <vector>
#include <algorithm>  // For std::copy

using namespace std;
using namespace Eigen;

vector<Vector3d> get_other_positions(size_t i, vector<Vector3d> positions) {
    vector<Vector3d> other_positions;
    other_positions.reserve(positions.size() - 1);  // 预留空间
    std::copy(positions.begin(), positions.begin() + i, std::back_inserter(other_positions));
    std::copy(positions.begin() + i + 1, positions.end(), std::back_inserter(other_positions));
    return other_positions;
}

// 判断所有飞机之间的三维距离是否都大于设定的阈值
bool are_all_planes_close(const vector<Vector3d>& positions, double threshold) {
    size_t num_robots = positions.size();
    
    // 遍历所有飞机的每一对
    for (size_t i = 0; i < num_robots; ++i) {
        for (size_t j = i + 1; j < num_robots; ++j) {
            double distance = (positions[i] - positions[j]).norm();
            if (distance <= threshold) {
                // 如果有任何一对飞机之间的距离小于等于阈值，返回true
                return true;
            }
        }
    }

    // 如果所有飞机之间的距离都大于阈值，返回false
    return false;
}

// 检查所有飞机是否与任何障碍物距离过近
bool are_planes_close_to_obstacles(const vector<Vector3d>& positions, const vector<Obstacle>& obstacles, double threshold = 1.0) {
    for (const auto& position : positions) {
        for (const auto& obstacle : obstacles) {
            double distance = (position - obstacle.position).norm();
            if (distance < obstacle.obstacle_radius + threshold) {
                // 如果飞机与障碍物的距离小于障碍物半径加上阈值，认为飞机与障碍物过近
                return true;
            }
        }
    }
    return false;
}


void generate_positions_in_grid(vector<Vector3d>& positions, size_t num_robots, double grid_size, double min_distance) {
    positions.clear();
    size_t grid_size_x = static_cast<size_t>(sqrt(num_robots));  // 计算网格的行数
    size_t grid_size_y = static_cast<size_t>(num_robots / grid_size_x);  // 计算网格的列数

    // 如果不是正方形的网格，调整列数
    if (grid_size_x * grid_size_y < num_robots) {
        grid_size_y++;
    }

    // 为每架飞机生成位置，按网格分布
    size_t count = 0;
    for (size_t i = 0; i < grid_size_x; ++i) {
        for (size_t j = 0; j < grid_size_y; ++j) {
            if (count >= num_robots) break;
            Vector3d new_pos(i * grid_size, j * grid_size, 0);  // 使用网格大小计算位置
            positions.push_back(new_pos);
            count++;
        }
    }
}




int main() {
    int num_robots = 128;  // 无人机数量
    vector<swarm_controller> controllers(num_robots);
    vector<Vector3d> positions(num_robots), other_positions(num_robots - 1), velocities(num_robots), other_velocities(num_robots - 1);
    Vector3d current_position, velocity;

    Vector3d goal(240, 240, 20);  // 目标点

    vector<Obstacle> obstacles;
    obstacles.emplace_back(Obstacle(Vector3d(90.0, 90.0, 0.0), 40.0));  // 添加障碍物
    obstacles.emplace_back(Obstacle(Vector3d(200.0, 150.0, 0.0), 20.0));  // 添加障碍物
    obstacles.emplace_back(Obstacle(Vector3d(170.0, 200.0, 0.0), 20.0));  // 添加障碍物

    // 使用网格布局生成无人机位置
    generate_positions_in_grid(positions, num_robots, 3, 2.2);  // 10.0是网格大小

    // 初始化每个无人机控制器
    for (int i = 0; i < num_robots; i++) {
        current_position = positions[i];
        other_positions = get_other_positions(i, positions);
        controllers[i].set_positions(current_position, other_positions);
        controllers[i].set_goal(goal);
        controllers[i].set_obstacles(obstacles);
    }

    double dT = 0.01;

    // 仿真20秒
    for (double t = 0; t < 40; t += dT) {
        for (int i = 0; i < num_robots; i++) {
            Vector3d speed = controllers[i].compute_force();
            positions[i] += speed * dT;  // 更新无人机位置
            std::cout << "position " << i << ": " << positions[i].transpose() << std::endl;
            current_position = positions[i];
            other_positions = get_other_positions(i, positions);
            controllers[i].set_positions(current_position, other_positions);
            velocities[i] = speed;
            other_velocities = get_other_positions(i, velocities);
        }

        // 检查飞机是否与障碍物过近
        if (are_planes_close_to_obstacles(positions, obstacles, 1.5)) {
            std::cout << "crash with obstacle!" << std::endl;
            break;  // 停止仿真，发生碰撞
        }

        // 检查飞机是否相距过近
        if (are_all_planes_close(positions, 2.0)) {
            std::cout << "crash with airplane!" << std::endl;
            break;  // 停止仿真，发生碰撞
        }
    }

    return 0;
}
